#include "llvm/IR/Function.h"
#include "llvm/IR/IRBuilder.h"
#include "llvm/IR/Constants.h"
#include "llvm/IR/Instructions.h"
#include "llvm/Pass.h"
#include "llvm/IR/Module.h"
#include "llvm/IR/BasicBlock.h"
#include "llvm/IR/Intrinsics.h"
#include "llvm/Support/raw_ostream.h"
#include "llvm/Transforms/Utils/Cloning.h"
#include "llvm/Transforms/Utils/ValueMapper.h"
#include "llvm/ADT/StringRef.h"
#include "llvm/ADT/Twine.h"
#include "llvm-c/Core.h"
#include "llvm/ADT/APInt.h"
#include "llvm/IR/GlobalVariable.h"
#include "llvm/IR/Type.h"
#include "llvm/Transforms/Utils/BasicBlockUtils.h"
#include "llvm/IR/LegacyPassManager.h"
#include "llvm/Transforms/IPO/PassManagerBuilder.h"
#include "Util.h"

#include <vector>
#include <unordered_set>
#include <random>
#include<iostream>
#include <algorithm>
#include <cstdint>
#include <queue>
#include <string>
#include<z3++.h>
#include <stdlib.h>
#include <time.h>

using namespace llvm;
using namespace z3;

namespace {
  class ObfuscateMBA : public FunctionPass {

    public:
    static char ID;
    ObfuscateMBA() : FunctionPass(ID) {}
    bool runOnFunction(Function &F) override;

    private:

    void MBA_generate(BinaryOperator* inst);
    uint64_t *getZeroVector(int logicLevels);
    Value* changeInst(uint64_t *value, BinaryOperator* inst);
  };
}

char ObfuscateMBA::ID = 0;
static RegisterPass<ObfuscateMBA> X("obfMBA", "obfuscate oprands with mba expresssions");
Pass *createObfuscateMBAPass() { return new ObfuscateMBA(); }

bool ObfuscateMBA::runOnFunction(Function &F) {
  for (BasicBlock &BB :F)
  {
      std::vector<Instruction*> IV;
      for (Instruction &I :BB)
      {
          IV.push_back(&I);
      }
      for (Instruction *I :IV)
      {
          if (isa<BinaryOperator>(I) && I->getOperand(0)->getType()->isIntegerTy())
          {
              BinaryOperator* inst = cast<BinaryOperator>(I);
              MBA_generate(inst);
          }
      }
  }
  return 1;
}

Value* ObfuscateMBA::changeInst(uint64_t *value, BinaryOperator* inst)
{
    IRBuilder<> builder(inst->getContext());
    builder.SetInsertPoint(inst);
    Value* a = NULL;
    Value* b = NULL;

    a = inst->getOperand(0);
    b = inst->getOperand(1);
       
    Value* expr = ConstantInt::get(a->getType(),0);
    Value* tmp = NULL;
    if (value[0] != 0)
    {
        tmp = builder.CreateAnd(a, b);
        tmp = builder.CreateMul(ConstantInt::get(a->getType(), value[0]), tmp);
        expr = builder.CreateAdd(expr, tmp);
    }
    if (value[1] != 0)
    {
        tmp = builder.CreateAnd(a, builder.CreateNot(b));
        tmp = builder.CreateMul(ConstantInt::get(a->getType(), value[1]), tmp);
        expr = builder.CreateAdd(expr, tmp);
    }
    if (value[2] != 0)
    {
        tmp = a;
        tmp = builder.CreateMul(ConstantInt::get(a->getType(), value[2]), tmp);
        expr = builder.CreateAdd(expr, tmp);
    }
    if (value[3] != 0)
    {
        tmp = builder.CreateAnd(builder.CreateNot(a), b);
        tmp = builder.CreateMul(ConstantInt::get(a->getType(), value[3]), tmp);
        expr = builder.CreateAdd(expr, tmp);
    }
    if (value[4] != 0)
    {
        tmp = b;
        tmp = builder.CreateMul(ConstantInt::get(a->getType(), value[4]), tmp);
        expr = builder.CreateAdd(expr, tmp);
    }
    if (value[5] != 0)
    {
        tmp = builder.CreateXor(a, b);
        tmp = builder.CreateMul(ConstantInt::get(a->getType(), value[5]), tmp);
        expr = builder.CreateAdd(expr, tmp);
    }
    if (value[6] != 0)
    {
        tmp = builder.CreateOr(a, b);
        tmp = builder.CreateMul(ConstantInt::get(a->getType(), value[6]), tmp);
        expr = builder.CreateAdd(expr, tmp);
    }
    if (value[7] != 0)
    {
        tmp = builder.CreateNot(builder.CreateOr(a,b));
        tmp = builder.CreateMul(ConstantInt::get(a->getType(), value[7]), tmp);
        expr = builder.CreateAdd(expr, tmp);
    }
    if (value[8] != 0)
    {
        tmp = builder.CreateNot(builder.CreateXor(a,b));
        tmp = builder.CreateMul(ConstantInt::get(a->getType(), value[8]), tmp);
        expr = builder.CreateAdd(expr, tmp);
    }
    if (value[9] != 0)
    {
        tmp = builder.CreateNot(b);
        tmp = builder.CreateMul(ConstantInt::get(a->getType(), value[9]), tmp);
        expr = builder.CreateAdd(expr, tmp);
    }
    if (value[10] != 0)
    {
        tmp = builder.CreateOr(builder.CreateNot(b), a);
        tmp = builder.CreateMul(ConstantInt::get(a->getType(), value[10]), tmp);
        expr = builder.CreateAdd(expr, tmp);
    }
    if (value[11] != 0)
    {
        tmp = builder.CreateNot(a);
        tmp = builder.CreateMul(ConstantInt::get(a->getType(), value[11]), tmp);
        expr = builder.CreateAdd(expr, tmp);
    }
    if (value[12] != 0)
    {
        tmp = builder.CreateOr(b, builder.CreateNot(a));
        tmp = builder.CreateMul(ConstantInt::get(a->getType(), value[12]), tmp);
        expr = builder.CreateAdd(expr, tmp);
    }
    if (value[13] != 0)
    {
        tmp = builder.CreateNot(builder.CreateAnd(a, b));
        tmp = builder.CreateMul(ConstantInt::get(a->getType(), value[13]), tmp);
        expr = builder.CreateAdd(expr, tmp);
    }
    if (value[14] != 0)
    {
        tmp = ConstantInt::get(a->getType(), -1*value[14]);
        expr = builder.CreateAdd(expr, tmp);
    }
    
    return expr;
}

void ObfuscateMBA::MBA_generate(BinaryOperator* inst)
{
    uint64_t * value = NULL;
    Value * expr = NULL;
    switch (inst->getOpcode()) {
        case BinaryOperator::Add:
            value = getZeroVector(5);
            value[2] += 1;
            value[4] += 1;
            expr = changeInst(value, inst);
            if (expr) inst->replaceAllUsesWith(expr);
            delete[] value;
            break;
        case BinaryOperator::Sub:
            value = getZeroVector(5);
            value[2] += 1;
            value[4] -= 1;
            changeInst(value, inst);
            expr = changeInst(value, inst);
            if (expr) inst->replaceAllUsesWith(expr);
            delete[] value;
            break;
        case BinaryOperator::And:
            value = getZeroVector(5);
            value[0] += 1;
            changeInst(value, inst);
            expr = changeInst(value, inst);
            if (expr) inst->replaceAllUsesWith(expr);
            delete[] value;
            break;
        case BinaryOperator::Or:
            value = getZeroVector(5);
            value[6] += 1;
            changeInst(value, inst);
            expr = changeInst(value, inst);
            if (expr) inst->replaceAllUsesWith(expr);
            delete[] value;
            break;
            
        
        default:
            break;
    }
    return;
}


static int8_t boolExprTable[15][4] = {
    {0, 0, 0, 1},
    {0, 0, 1, 0},
    {0, 0, 1, 1},
    {0, 1, 0, 0},
    {0, 1, 0, 1},
    {0, 1, 1, 0},
    {0, 1, 1, 1},
    {1, 0, 0, 0},
    {1, 0, 0, 1},
    {1, 0, 1, 0},
    {1, 0, 1, 1},
    {1, 1, 0, 0},
    {1, 1, 0, 1},
    {1, 1, 1, 0},
    {1, 1, 1, 1}
};

uint64_t *ObfuscateMBA::getZeroVector(int logicLevels)
{
    srand(time(NULL));
    int *exprs = new int[logicLevels];
    uint64_t *coefficients = new uint64_t[15];
    while (1)
    {
        context c;
        std::vector<expr> element;
        solver s(c);
        std::fill_n(coefficients, 15, 0);
        for (int i=0; i <logicLevels; i++)
        {
            std::string name;
            name += 'x';
            name += std::to_string(i);
            element.push_back(c.int_const(name.c_str()));
        }
        for (int i=0; i<logicLevels;i++)
        {
            exprs[i] = rand()%15;
        }
        for (int i=0;i<4; i++)
        {
            expr tmpExpr = c.int_val(0);
            for (int j=0;j<logicLevels; j++)
            {
                tmpExpr = tmpExpr+element[j]*boolExprTable[exprs[j]][i];
            }
            s.add( tmpExpr == 0);
        }
        expr constraints = c.bool_val(false);
        for (int i = 0; i < logicLevels; i++) {
            constraints = constraints || (element[i] != 0);
        }
        s.add(constraints);
        if (s.check() != sat)
        {
            errs() << "z3 solve failed!" <<"\n";
            continue;
        }
        model m = s.get_model();
        for (int i = 0; i < logicLevels; i++)
        {
            coefficients[exprs[i]] += m.eval(element[i]).get_numeral_int64();
        }
        delete[] exprs;
        return coefficients;
    }
}
